"""
工作流基础模块
定义工作流基础类和工具函数
"""
from typing import Dict, List, Any, Optional, Callable, TypeVar, Generic
from loguru import logger
import datetime
import json
from pathlib import Path

T = TypeVar('T')

class WorkflowState(Generic[T]):
    """工作流状态类"""
    
    def __init__(self, initial_state: Optional[T] = None):
        """
        初始化工作流状态
        
        参数:
            initial_state: 初始状态
        """
        self.state = initial_state or {}
        self.history = []
        self.created_at = datetime.datetime.now()
        self.last_updated = self.created_at
    
    def update(self, new_state: T) -> None:
        """
        更新状态
        
        参数:
            new_state: 新状态
        """
        # 记录历史
        self.history.append((self.state, self.last_updated))
        
        # 更新状态
        self.state = new_state
        self.last_updated = datetime.datetime.now()
    
    def get_state(self) -> T:
        """
        获取当前状态
        
        返回:
            当前状态
        """
        return self.state
    
    def get_history(self) -> List[tuple]:
        """
        获取状态历史
        
        返回:
            状态历史列表
        """
        return self.history
    
    def save(self, filepath: str) -> None:
        """
        保存状态到文件
        
        参数:
            filepath: 文件路径
        """
        try:
            # 确保目录存在
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)
            
            # 准备保存的数据
            data = {
                "state": self.state,
                "created_at": self.created_at.isoformat(),
                "last_updated": self.last_updated.isoformat(),
                "history": [(state, timestamp.isoformat()) for state, timestamp in self.history]
            }
            
            # 保存到文件
            with open(filepath, 'w') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            logger.info(f"工作流状态已保存到 {filepath}")
        
        except Exception as e:
            logger.error(f"保存工作流状态失败: {e}")
    
    @classmethod
    def load(cls, filepath: str) -> 'WorkflowState':
        """
        从文件加载状态
        
        参数:
            filepath: 文件路径
            
        返回:
            工作流状态实例
        """
        try:
            # 检查文件是否存在
            if not Path(filepath).exists():
                logger.warning(f"状态文件 {filepath} 不存在，返回新状态")
                return cls()
            
            # 从文件加载
            with open(filepath, 'r') as f:
                data = json.load(f)
            
            # 创建实例
            instance = cls(data["state"])
            instance.created_at = datetime.datetime.fromisoformat(data["created_at"])
            instance.last_updated = datetime.datetime.fromisoformat(data["last_updated"])
            instance.history = [(state, datetime.datetime.fromisoformat(timestamp)) 
                              for state, timestamp in data["history"]]
            
            logger.info(f"已从 {filepath} 加载工作流状态")
            return instance
        
        except Exception as e:
            logger.error(f"加载工作流状态失败: {e}")
            return cls()


class BaseWorkflow:
    """工作流基类"""
    
    def __init__(self, name: str, state: Optional[WorkflowState] = None):
        """
        初始化工作流
        
        参数:
            name: 工作流名称
            state: 工作流状态
        """
        self.name = name
        self.state = state or WorkflowState({})
        self.steps = []
        self.current_step_index = 0
        
        logger.info(f"初始化工作流: {name}")
    
    def add_step(self, step_func: Callable, name: str) -> None:
        """
        添加工作流步骤
        
        参数:
            step_func: 步骤函数
            name: 步骤名称
        """
        self.steps.append({"func": step_func, "name": name})
        logger.info(f"添加工作流步骤: {name}")
    
    def run(self) -> Any:
        """
        运行工作流
        
        返回:
            工作流运行结果
        """
        logger.info(f"开始运行工作流: {self.name}")
        
        result = None
        for i, step in enumerate(self.steps):
            self.current_step_index = i
            step_name = step["name"]
            step_func = step["func"]
            
            logger.info(f"运行步骤 {i+1}/{len(self.steps)}: {step_name}")
            try:
                result = step_func(self.state.get_state(), result)
                
                # 更新状态
                current_state = self.state.get_state()
                if isinstance(current_state, dict) and isinstance(result, dict):
                    # 合并字典
                    updated_state = {**current_state, **result}
                    self.state.update(updated_state)
                else:
                    # 直接替换
                    self.state.update(result)
                
                logger.info(f"步骤 {step_name} 完成")
            
            except Exception as e:
                logger.error(f"步骤 {step_name} 运行失败: {e}")
                import traceback
                logger.error(traceback.format_exc())
                break
        
        logger.info(f"工作流 {self.name} 运行完成")
        return self.state.get_state()
    
    def run_step(self, step_index: int, input_data: Any = None) -> Any:
        """
        运行指定步骤
        
        参数:
            step_index: 步骤索引
            input_data: 输入数据
            
        返回:
            步骤运行结果
        """
        if step_index < 0 or step_index >= len(self.steps):
            logger.error(f"步骤索引 {step_index} 超出范围")
            return None
        
        step = self.steps[step_index]
        step_name = step["name"]
        step_func = step["func"]
        
        logger.info(f"运行步骤: {step_name}")
        try:
            result = step_func(self.state.get_state(), input_data)
            logger.info(f"步骤 {step_name} 完成")
            return result
        
        except Exception as e:
            logger.error(f"步骤 {step_name} 运行失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return None
    
    def save_state(self, filepath: str) -> None:
        """
        保存工作流状态
        
        参数:
            filepath: 文件路径
        """
        self.state.save(filepath)
    
    def load_state(self, filepath: str) -> None:
        """
        加载工作流状态
        
        参数:
            filepath: 文件路径
        """
        self.state = WorkflowState.load(filepath)